from langchain_core.prompts import ChatPromptTemplate
from customize.get_ollama import GetOllama
from pydantic import BaseModel, Field
from langchain_core.output_parsers import PydanticOutputParser
import pandas as pd
from langchain_community.document_loaders import PythonLoader
import os

prompt_template = ChatPromptTemplate.from_template(
    """你是一位高中的python教师，你的学生是python的初学者，你正在批改学生python程序的作业，作业共分为A，B，C，D，E五个等级。
请你根据以下题目内容、作业初始状态(有可能空白)、作业的参考答案对学生提交的作业进行综合评价，评价可以从输入程序是否正确，格式是否规范，运行结果是否正确等方面综合考虑，
最终由高至低给出A~E的评价,并用简洁的语言,以老师对学生说话的方式，说明给出等级的理由，如果错误，指出具体错误。
题目:\n
{question}
初始的作业：\n
{original}
参考答案：\n
{right_answer}
学生提交作业：\n
{assignments}

请尝试从作业文本的开头，提取学生的基本信息，包括班级，姓名，学号，如果提取不到，请尝试从文件名：{file_name}进行推断，如果推断不出来，请请留空白。
请使用以下的格式返回结果：
    "class_num":"学生班级"
    "student":"学生姓名"
    "id":"学号"
    "grade":"学生等级",
    "evaluation"："给学生的评级"
输出时评价时，对特殊符号进行转义处理，转义符号使用\\, 避免使用`作为转义符，保证输出的的值不会导致python程序出现字符串错误。
"""

)


class StructResult(BaseModel):
    class_num: str = Field(description="班级")
    student: str = Field(description="学生姓名")
    id: str = Field(description="学号")
    grade: str = Field(description="学生作业等级", example="A")
    evaluation: str = Field(description="给学生的一段综合的评价")


# llm = GetOllama(model_name="qwen2.5-coder:7b", model_type=1)()
# question = """
# 题目：学生考试成绩统计
# 假设有 10 位同学的数学考试成绩如下：
# scores = [85, 90, 78, 88, 92, 75, 80, 95, 86, 70]
# 请编写一个 Python 程序，实现以下功能：
# 1.使用循环计算这 10 位同学的平均成绩。
# 2.找出成绩高于平均成绩的同学的分数，并打印出来。
# """
# original = """
# # 班级：
# # 姓名：
# # 学号：
#
# scores = [85, 90, 78, 88, 92, 75, 80, 95, 86, 70]
#
# # 总分变量
# total_score =
# # 累计总分
# for score in scores:
#
#
# # 计算平均成绩（使用len()计算列表元素个数）
# average_score =
#
# # 找出高于平均成绩的同学的分数
#
#
#         print(f"{score} 超过了平均分{average_score}。")
# """
# right_answer = """
# # 班级：
# # 姓名：
# # 学号：
#
# scores = [85, 90, 78, 88, 92, 75, 80, 95, 86, 70]
#
# # 总分变量
# total_score = 0
# # 累计总分
# for score in scores:
#     total_score = total_score + score
#
# # 计算平均成绩（使用len()计算列表元素个数）
# average_score = total_score / len(scores)
#
# # 找出高于平均成绩的同学的分数
# for score in scores:
#     if score > average_score:
#         print(f"{score} 超过了平均分{average_score}。")
# """
# assignments = """
# # 班级：108
# # 姓名：孟思彤
# # 学号：34
#
# scores = [85, 90, 78, 88, 92, 75, 80, 95, 86, 70]
#
# # 总分变量
# total_score = 0
# # 累计总分
# for score in scores:
#     total_score=score+total_score
# # 计算平均成绩（使用len()计算列表元素个数）
# average_score =total_score/len(scores)
#
# # 找出高于平均成绩的同学的分数
# for score in scores:
#     if(score>average_score):
#         print(f"{score} 超过了平均分{average_score}。")
# """

# prompt = prompt_template.fromat(
#     question=question,
#     original=original,
#     correct_answer=right_answer,
#     assignments=assignments
# )


# parser = PydanticOutputParser(pydantic_object=StructResult)
# chain = prompt_template | llm | parser


# response = chain.invoke({'question': question, 'original': original, 'right_answer': right_answer,
# 'assignments': assignments})
# print(response)

def get_python_files_in_folders(folder_path):
    python_files = []
    for root, dirs, files in os.walk(folder_path):
        for file in files:
            if file.endswith('.py'):
                full_path = os.path.join(root, file)
                python_files.append(full_path)

    return python_files


class ExerciseCorrection:
    def __init__(self, folder, question, right_answer, original=""):
        self.folder_path = folder
        self.llm = GetOllama(model_name="qwen2.5-coder:7b", model_type=1)()
        self.question = question
        self.original = original
        self.right_answer = right_answer
        self.prompt_template = prompt_template
        self.parser = PydanticOutputParser(pydantic_object=StructResult)
        self.chain = self.prompt_template | self.llm | self.parser
        self.handle_model_response = None

    def get_python_files_in_folders(self):
        python_files = []
        for root, dirs, files in os.walk(self.folder_path):
            for file in files:
                if file.endswith('.py'):
                    full_path = os.path.join(root, file)
                    python_files.append(full_path)

        return python_files

    def __call__(self):
        python_files = self.get_python_files_in_folders()
        results = []
        for python_file in python_files:
            loader = PythonLoader(python_file)
            documents = loader.load()
            response = self.chain.invoke(
                {
                    'question': self.question,
                    'original': self.original,
                    'right_answer': self.right_answer,
                    'assignments': documents[0].page_content,
                    'file_name': os.path.basename(python_file)
                }
            )
            if self.handle_model_response is not None:
                self.handle_model_response(response.dict())
            # print(response)
            results.append(response)

        # # 将对象列表转换为字典列表
        # data_dict = [d.dict() for d in results]
        #
        # # 使用pandas转换为DataFrame并保存为Excel
        # df = pd.DataFrame(data_dict)
        return True
        # df.to_excel(save_path, index=False)

# if __name__ == '__main__':
#     folder_path = r"J:\2024学生python作业\8班"
#     save_path = "result.xlsx"
#     python_files = get_python_files_in_folders(folder_path)
#     results = []
#     for python_file in python_files:
#         loader = PythonLoader(python_file)
#         documents = loader.load()
#         response = chain.invoke({'question': question, 'original': original, 'right_answer': right_answer,
#                                  'assignments': documents[0].page_content})
#         print(response)
#         results.append(response)
#
#     # 将对象列表转换为字典列表
#     data_dict = [d.dict() for d in results]
#
#     # 使用pandas转换为DataFrame并保存为Excel
#     df = pd.DataFrame(data_dict)
#     df.to_excel(save_path, index=False)